
#include "default_defines.h"
#include "global_definitions.h"
#include "device.h"
#include "openmp.h"
#include "sotl.h"

#ifdef HAVE_LIBGL
#include "vbo.h"
#endif

#include <stdio.h>
#include <string.h>
#include <omp.h>
#include <math.h>
#include <assert.h>

static int *atom_state = NULL;

// Used to color threads
static struct {
  float R,G,B;
} proc_color [] = {
  {1.0, 0.0, 0.0},
  {1.0, 1.0, 0.0},
  {1.0, 0.0, 1.0},
  {1.0, 1.0, 0.0},
  {0.0, 1.0, 1.0},
  {0.0, 0.0, 1.0},
  {1.0, 0.5, 0.0},
  {1.0, 0.0, 0.5},
  {0.5, 1.0, 0.0},
  {0.0, 1.0, 0.5},
  {0.5, 0.0, 1.0},
  {0.0, 0.5, 1.0},
};

#ifdef HAVE_LIBGL

#define SHOCK_PERIOD  50

// Update OpenGL Vertex Buffer Object
//
static void omp_update_vbo (sotl_device_t *dev)
{
  sotl_atom_set_t *set = &dev->atom_set;
  //sotl_domain_t *domain = &dev->domain;

#pragma omp parallel for schedule(runtime)
    for (unsigned n = 0; n < set->natoms; n++) {
      vbo_vertex[n*3 + 0] = set->pos.x[n];
      vbo_vertex[n*3 + 1] = set->pos.y[n];
      vbo_vertex[n*3 + 2] = set->pos.z[n];
      
      /*// Atom color depends on z coordinate
	{
	float ratio = (set->pos.z[n] - domain->min_ext[2]) / (domain->max_ext[2] - domain->min_ext[2]);
	
	vbo_color[n*3 + 0] = (1.0 - ratio) * atom_color[0].R + ratio * 1.0;
	vbo_color[n*3 + 1] = (1.0 - ratio) * atom_color[0].G + ratio * 0.0;
	vbo_color[n*3 + 2] = (1.0 - ratio) * atom_color[0].B + ratio * 0.0;
	atom_state[n]--;
	}*/
      
      /*// Atom color depends on bounces
      {
	float ratio = ((float)atom_state[n]) / ((float)SHOCK_PERIOD);
	
	vbo_color[n*3 + 0] = (1.0 - ratio) * atom_color[0].R + ratio * 1.0;
	vbo_color[n*3 + 1] = (1.0 - ratio) * atom_color[0].G + ratio * 0.0;
	vbo_color[n*3 + 2] = (1.0 - ratio) * atom_color[0].B + ratio * 0.0;
	atom_state[n]--;
    }*/

      // Atom color depends on threads and on bounces
      {
	float ratio = ((float)atom_state[n]) / ((float)SHOCK_PERIOD);
	
	vbo_color[n*3 + 0] = (1.0 - ratio) * proc_color[omp_get_thread_num()].R + ratio * 1.0;
	vbo_color[n*3 + 1] = (1.0 - ratio) * proc_color[omp_get_thread_num()].G + ratio * 0.0;
	vbo_color[n*3 + 2] = (1.0 - ratio) * proc_color[omp_get_thread_num()].B + ratio * 0.0;
	atom_state[n]--;
    }
  }
}
#endif

// Update positions of atoms by adding (dx, dy, dz)
//
static void omp_move (sotl_device_t *dev)
{
  sotl_atom_set_t *set = &dev->atom_set;
  sotl_domain_t *domain = &dev->domain;

#pragma omp parallel for schedule(runtime)
  for (unsigned n = 0; n < set->natoms; n++) {
    set->pos.x[n] += set->speed.dx[n];
    set->pos.y[n] += set->speed.dy[n];
    set->pos.z[n] += set->speed.dz[n];

    //With the lennard potential, atoms can be thrown far away. To prevent that, if an atom goes out of the domain, we make it re-enter by the opposite side
    for(unsigned coord = 0; coord < 3; ++coord)
      {
	double pos = set->pos.x[coord * set->offset + n];
	double min_pos = domain->min_ext[coord];
	double max_pos = domain->max_ext[coord];

	//we put the atom back into the domain using modulus
	if(pos-min_pos < 0)
	  set->pos.x[coord * set->offset + n] = max_pos - fmod(-1 * (pos-min_pos), max_pos - min_pos);
	else
	  set->pos.x[coord * set->offset + n] = fmod(pos, max_pos - min_pos) + min_pos;
      }
  }
}

// Apply gravity force
//
static void omp_gravity (sotl_device_t *dev)
{
  /*
  sotl_atom_set_t *set = &dev->atom_set;
  const calc_t g = 0.005;

  //TODO
  */
}

static void omp_bounce (sotl_device_t *dev)
{
  sotl_atom_set_t *set = &dev->atom_set;
  sotl_domain_t *domain = &dev->domain;


#pragma omp parallel for schedule(runtime)
  for(unsigned n = 0; n < set->natoms; n++) {
    for(unsigned j = 0; j < 3; j++) {
      if((set->pos.x[n + j * set->offset] <= domain->min_ext[j]) || (set->pos.x[n + j * set->offset] >= domain->max_ext[j])) {
	set->speed.dx[n + j * set->offset] *= -1;
	
	// Version with friction
	//set->speed.dx[n + j * set->offset] *= -0.9;
	
	atom_state[n] = SHOCK_PERIOD;
      }
    }
  }
}

static calc_t squared_distance (sotl_atom_set_t *set, unsigned p1, unsigned p2)
{
  calc_t *pos1 = set->pos.x + p1,
    *pos2 = set->pos.x + p2;

  calc_t dx = pos2[0] - pos1[0],
         dy = pos2[set->offset] - pos1[set->offset],
         dz = pos2[set->offset*2] - pos1[set->offset*2];

  return dx * dx + dy * dy + dz * dz;
}

static calc_t lennard_jones (calc_t r2)
{
  calc_t rr2 = 1.0 / r2;
  calc_t r6;

  r6 = LENNARD_SIGMA * LENNARD_SIGMA * rr2;
  r6 = r6 * r6 * r6;

  return 24 * LENNARD_EPSILON * rr2 * (2.0f * r6 * r6 - r6);
}

//Auxiliar fonction which computes, for each atom in box1, the interaction of each atom of box2 on it. This fonction isn't symmetric, it only computes forces from box2 over box1.
static void omp_force_between_two_different_boxes(sotl_atom_set_t* sorted_set, int boxes[], unsigned int box1, unsigned int box2)
{
  //current1 is the index of the atom of box1. Same for current2 and box2.

  for(unsigned int current1 = boxes[box1]; (int)current1 < boxes[box1 + 1]; ++current1)
    {
      //iteration over the second box
      for(unsigned int current2 = boxes[box2]; (int)current2 < boxes[box2 + 1]; ++current2)
	{
	  
	  calc_t sq_dist = squared_distance(sorted_set, current1, current2);

	  if(sq_dist < LENNARD_SQUARED_CUTOFF)
	    {

	      calc_t intensity = lennard_jones(sq_dist);
	      calc_t force[3] = {0.0, 0.0, 0.0};
	      
	      force[0] += intensity * (sorted_set->pos.x[current1] - sorted_set->pos.x[current2]);
	      force[1] += intensity * (sorted_set->pos.x[sorted_set->offset + current1] -
				       sorted_set->pos.x[sorted_set->offset + current2]);
	      force[2] += intensity * (sorted_set->pos.x[sorted_set->offset * 2 + current1] -
				       sorted_set->pos.x[sorted_set->offset * 2 + current2]);
	      
	      sorted_set->speed.dx[current1] += force[0];
	      sorted_set->speed.dx[sorted_set->offset + current1] += force[1];
	      sorted_set->speed.dx[sorted_set->offset * 2 + current1] += force[2];
	    }
	}
      
    }
}

//Auxiliar fonction which computes the interactions between all the atoms of the same box.
static void omp_force_in_one_box(sotl_atom_set_t* sorted_set, int boxes[], unsigned int box)
{

  for(unsigned int current = boxes[box]; (int)current < boxes[box + 1]; ++current)
    {

      for(unsigned int other = boxes[box]; (int)other < boxes[box + 1]; ++other)
	{

	  //that if-statement is the only difference with omp_force_between_two_different_boxes()
	  if(current != other)
	    {
	      
	      calc_t sq_dist = squared_distance(sorted_set, current, other);

	      if(sq_dist < LENNARD_SQUARED_CUTOFF)
		{
		  calc_t intensity = lennard_jones(sq_dist);
		  calc_t force[3] = {0.0, 0.0, 0.0};
		  
		  force[0] += intensity * (sorted_set->pos.x[current] - sorted_set->pos.x[other]);
		  force[1] += intensity * (sorted_set->pos.x[sorted_set->offset + current] -
					   sorted_set->pos.x[sorted_set->offset + other]);
		  force[2] += intensity * (sorted_set->pos.x[sorted_set->offset * 2 + current] -
					   sorted_set->pos.x[sorted_set->offset * 2 + other]);
		  
		  sorted_set->speed.dx[current] += force[0];
		  sorted_set->speed.dx[sorted_set->offset + current] += force[1];
		  sorted_set->speed.dx[sorted_set->offset * 2 + current] += force[2];
		}
	    }
	}
    }
}

//Version with boxes
static void omp_force (sotl_device_t *dev)
{
  sotl_atom_set_t *set = &dev->atom_set;
  sotl_domain_t *domain = &dev->domain;


  //Sorts the atoms


  //Computes the number of atom in each box, and stores it in a vector

  int* boxes = atom_set_box_count(domain, set);

  //Computes, for each box, the index of the 1st atom in the box. Stores it in boxes[]

  unsigned int atoms_without_box = set->natoms;
  for(int j = domain->total_boxes - 1; j >= 0; --j) {
    unsigned int mem = atoms_without_box - boxes[j];
    boxes[j] = mem;
    atoms_without_box = mem;
  }

//A vector containing, for each box, the index of the next empty slot in the new sorted vector. Used to fill the vector progressively wih the atoms
  int* box_next_empty = malloc(domain->total_boxes * sizeof(int));

//At the beginning, "box_next_empty" is equal to "boxes"
  memcpy(box_next_empty, boxes, domain->total_boxes*sizeof(int));

  int * boxes_tmp3 = atom_set_box_count(domain, set);

  //We copy each atom at the right position in the copied vectors
  for(unsigned int i = 0; i < set->natoms; ++i) {

    int box_number = atom_get_num_box(domain, set->pos.x[i], set->pos.y[i], set->pos.z[i], BOX_SIZE_INV);

    //Copies the atom in its new place in set->sorted_pos and set->sorted_speed

    int new_index = box_next_empty[box_number];
    
    set->sorted_pos.x[new_index] = set->pos.x[i];
    set->sorted_pos.x[new_index + set->offset] = set->pos.y[i];
    set->sorted_pos.x[new_index + 2 * set->offset] = set->pos.z[i];
    
    set->sorted_speed.dx[new_index] = set->speed.dx[i];
    set->sorted_speed.dx[new_index + set->offset] = set->speed.dy[i];
    set->sorted_speed.dx[new_index + 2 * set->offset] = set->speed.dz[i];

    ++(box_next_empty[box_number]);
  }

  //Replacement of the old vectors by the new sorted ones
  memcpy(set->pos.x, set->sorted_pos.x, 3 * set->offset * sizeof(calc_t));
  memcpy(set->speed.dx, set->sorted_speed.dx, 3 * set->offset * sizeof(calc_t));

  //
  //Computes the forces and applies it
  //

  //We iterate through the boxes, and for each box we compute all the interactions due to the surrounding boxes.

  //We start with the first "real" box, i.e. the 1st box which is not a border one. Assumming the box are indexed from (x=0, y=0, z=0), this correspond to box (x=1, y=1, z=1)

  unsigned int start = 1 + domain->boxes[0] + domain->boxes[0] * domain->boxes[1];

  //In the same way, we end with the last "real" box ( i.e. box (xmax-2, ymax-2, zmax-2) )

  unsigned int end = (domain->boxes[0] - 2) + (domain->boxes[0] * (domain->boxes[1] - 2)) + (domain->boxes[0] * domain->boxes[1] * (domain->boxes[2] - 2));

  #pragma omp parallel for schedule(runtime)
  for(unsigned int current_box = start; current_box <= end; ++current_box) {

    //
    //z-1
    //


    //y-1
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box - domain->boxes[0] * domain->boxes[1] - domain->boxes[0] - 1); //x-1
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box - domain->boxes[0] * domain->boxes[1] - domain->boxes[0]); //x
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box - domain->boxes[0] * domain->boxes[1] - domain->boxes[0] + 1); //x+1

    //y
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box - domain->boxes[0] * domain->boxes[1] - 1); //x-1
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box - domain->boxes[0] * domain->boxes[1]); //x
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box - domain->boxes[0] * domain->boxes[1] + 1); //x+1

    //y+1
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box - domain->boxes[0] * domain->boxes[1] + domain->boxes[0] - 1); //x-1
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box - domain->boxes[0] * domain->boxes[1] + domain->boxes[0]); //x
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box - domain->boxes[0] * domain->boxes[1] + domain->boxes[0] + 1); //x+1


    //
    //z
    //


    //y-1
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box - domain->boxes[0] - 1); //x-1
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box - domain->boxes[0]); //x
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box - domain->boxes[0] * domain->boxes[1] - domain->boxes[0] + 1); //x+1

    //y
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box - 1); //x-1
    omp_force_in_one_box(set, boxes, current_box); //same box
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box + 1); //x+1

    //y+1
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box + domain->boxes[0] - 1); //x-1
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box + domain->boxes[0]); //x
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box + domain->boxes[0] + 1); //x+1


    //
    //z+1
    //


    //y-1
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box + domain->boxes[0] * domain->boxes[1] - domain->boxes[0] - 1); //x-1
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box + domain->boxes[0] * domain->boxes[1] - domain->boxes[0]); //x
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box + domain->boxes[0] * domain->boxes[1] - domain->boxes[0] + 1); //x+1

    //y
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box + domain->boxes[0] * domain->boxes[1] - 1); //x-1
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box + domain->boxes[0] * domain->boxes[1]); //x
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box + domain->boxes[0] * domain->boxes[1] + 1); //x+1

    //y+1
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box + domain->boxes[0] * domain->boxes[1] + domain->boxes[0] - 1); //x-1
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box + domain->boxes[0] * domain->boxes[1] + domain->boxes[0]); //x
    omp_force_between_two_different_boxes(set, boxes, current_box, current_box + domain->boxes[0] * domain->boxes[1] + domain->boxes[0] + 1); //x+1

  }
  
  //Frees the temporary vectors
  free(boxes);
  free(box_next_empty);

}



// Main simulation function
//
void omp_one_step_move (sotl_device_t *dev)
{

  // Apply gravity force
  //
  if (gravity_enabled)
    omp_gravity (dev);

  // Compute interactions between atoms
  //
  //if (force_enabled)
  omp_force (dev);
  
  //NOT TO USE
  // Bounce on borders
  //
  //if(borders_enabled)
  //omp_bounce (dev);
  
  // Update positions
  //
  omp_move (dev);

#ifdef HAVE_LIBGL
  // Update OpenGL position
  //
  if (dev->display)
    omp_update_vbo (dev);
#endif
}

void omp_init (sotl_device_t *dev)
{
#ifdef _SPHERE_MODE_
  sotl_log(ERROR, "Sequential implementation does currently not support SPHERE_MODE\n");
  exit (1);
#endif

  borders_enabled = 1;

  dev->compute = SOTL_COMPUTE_OMP; // dummy op to avoid warning
}

void omp_alloc_buffers (sotl_device_t *dev)
{
  atom_state = calloc(dev->atom_set.natoms, sizeof(int));
  printf("natoms: %d\n", dev->atom_set.natoms);
}

void omp_finalize (sotl_device_t *dev)
{
  free(atom_state);

  dev->compute = SOTL_COMPUTE_OMP; // dummy op to avoid warning
}






